from pydantic import BaseModel, Field
import typing as T


class Parameters(BaseModel):
    size: T.Optional[str] = Field(
        default="1024*1024",
        description="图像宽高边长的像素范围为：[768, 1440]",
    )
    n: T.Optional[int] = Field(
        default=1,
        description="生成图像的数量，默认为1，最大值为10。",
        ge=1,
        le=4,
    )
    seed: T.Optional[int] = Field(
        default=None,
        description="""随机数种子，用于控制模型生成内容的随机性。seed参数取值范围是[0, 2147483647]。如果不提供，则算法自动生成一个随机数作为种子。如果给定了，
        则根据n的值分别为n张图片生成seed参数，例如n=4，算法将分别生成seed、seed+1、seed+2、seed+3作为参数的图片。""".strip(),
        le=2147483647,
        ge=0,
    )
    prompt_extend: T.Optional[bool] = Field(
        default=False,
        description="""是否开启prompt智能改写。开启后会使用大模型对输入prompt进行智能改写，仅对正向提示词有效。
                        对于较短的输入prompt生成效果提升明显，但会增加3-4秒耗时。【默认为False】""".strip(),
    )
    watermark: T.Optional[bool] = Field(default=False, description="是否添加水印")


class Input(BaseModel):
    prompt: str = Field(
        description="正向提示词，用来描述生成图像中期望包含的元素和视觉特点。",
        max_length=800,
    )
    negative_prompt: T.Optional[str] = Field(
        default=None,
        description="反向提示词，用来描述生成图像中不希望包含的元素和视觉特点。",
        max_length=500,
    )


class GeneratorImageDTO(BaseModel):
    input: Input = Field(description="对图片的描述【prompt必填】")

    model: T.Literal["wanx2.1-t2i-turbo"] = Field(
        default="wanx2.1-t2i-turbo",
        description="模型名称【可选】目前只支持wanx2.1-t2i-turbo",
    )

    parameters: T.Optional[Parameters] = Field(
        default=Parameters(),
        description="图像生成参数【可选】",
    )
